'''
Created on Nov 4, 2009

@author: mkiyer
'''

'''
Created on Nov 4, 2009

@author: mkiyer
'''

import sys
from veggie.io.bed import simple_bed_reader, bed_reader

def convert_simple_bed(fhd, source, prefix):
    for linenum, annot in enumerate(simple_bed_reader(open(fhd))):
        gene_id = '%s_%d' % (prefix, linenum+1)
        transcript_id = gene_id + '.0'
        chrom, start, end = annot
        gtf_fields = [chrom,
                      source,
                      'transcript',
                      str(start),
                      str(end),
                      '0.0',
                      '.',
                      '.',
                      'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]
        print '\t'.join(gtf_fields)        
        gtf_fields = [chrom,
                      source,
                      'exon',
                      str(start),
                      str(end),
                      '0.0',
                      '.',
                      '.',
                      'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]        
        print '\t'.join(gtf_fields)   

def convert_bed(fhd, source, prefix=None):
    for linenum, annot in enumerate(bed_reader(open(sys.argv[1]))):
        if prefix is not None:
            gene_id = prefix + '.' + str(linenum+1)
        else:
            gene_id = annot.name
        transcript_id = gene_id + '.0'
        
        strand = annot.strand if annot.strand is not None else '+'
                
        if annot.blockCount == 0:
            gtf_fields = [annot.chrom,
                          source,
                          'exon',
                          str(annot.start),
                          str(annot.end),
                          '0.00',
                          strand,
                          '.',
                          'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]        
            print '\t'.join(gtf_fields)
        else:
            for length, offset in zip(annot.blockSizes, annot.blockStarts):
                gtf_fields = [annot.chrom,
                              source,
                              'exon',
                              str(annot.start + offset),
                              str(annot.start + offset + length),
                              '0.00',
                              strand,
                              '.',
                              'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]
                print '\t'.join(gtf_fields)

def convert_raw_file(fhd, source, prefix, chrom=1, start=2, end=3, name=None):
    for linenum, line in enumerate(fhd):
        line = line.strip()
        if line is None:
            continue
        if line.startswith('#'):
            continue
        fields = line.split('\t')
        if name is not None:
            n = fields[name]
        else:
            n = None
        convert_simple(fields[chrom], fields[start], fields[end], n, source, prefix, linenum)

def convert_simple(chrom, start, end, name, source, prefix, id):
    if name is not None:
        gene_id = '%s_%s_%d' % (prefix, name, id+1)
    else:
        gene_id = '%s_%d' % (prefix, id+1)
    transcript_id = gene_id + '.0'       
    gtf_fields = [chrom,
                  source,
                  'exon',
                  str(start),
                  str(end),
                  '0.000',
                  '+',
                  '.',
                  'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]        
    print '\t'.join(gtf_fields)  

if __name__ == '__main__':
    source = sys.argv[2]
    prefix = sys.argv[3]
    convert_raw_file(open(sys.argv[1], 'rU'), source, prefix, chrom=0, start=1, end=2, name=None)
    #convert_bed(open(sys.argv[1], 'rU'), source, prefix)
